package com.litian.dancechar.framework.common.trace;

import cn.hutool.core.util.StrUtil;
import com.litian.dancechar.framework.common.util.DCMd5Util;
import org.jboss.logging.MDC;
import org.springframework.core.NamedThreadLocal;

import java.util.UUID;

/**
 * 获取TraceId工具类
 *
 * @author litian
 * @date 2021/6/14 21:15
 */
public class TraceHelper {
    public static final ThreadLocal<Trace> CURRENT_SPAN = new NamedThreadLocal<>("TraceId Context");

    public static void setCurrentTrace(String traceId) {
        if (StrUtil.isBlank(traceId)) {
            traceId = DCMd5Util.getTraceId(UUID.randomUUID().toString());
        }
        Trace trace = new Trace();
        trace.setTraceId(traceId);
        trace.setSpanId(DCMd5Util.getTraceId(UUID.randomUUID().toString()));
        MDC.put(Trace.TRACE, trace.getTraceId());
        MDC.put(Trace.PARENT_SPAN, trace.getSpanId());
        CURRENT_SPAN.set(trace);
    }

    public static Trace getCurrentTrace() {
        Trace trace = CURRENT_SPAN.get();
        if (trace == null) {
            trace = new Trace();
            trace.setTraceId(DCMd5Util.getTraceId(UUID.randomUUID().toString()));
            trace.setSpanId(DCMd5Util.getTraceId(UUID.randomUUID().toString()));
            MDC.put(Trace.TRACE, trace.getTraceId());
            MDC.put(Trace.PARENT_SPAN, trace.getSpanId());
            CURRENT_SPAN.set(trace);
            return trace;
        }
        return trace;
    }

    public static String getTraceId() {
        return getCurrentTrace() != null ? getCurrentTrace().getTraceId() : null;
    }

    public static void clearCurrentTrace() {
        CURRENT_SPAN.set(null);
        CURRENT_SPAN.remove();
    }
}